from uuid import UUID
from datetime import datetime
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func
from sqlalchemy.orm import joinedload
from typing import List, Tuple, Optional

from app.crud.base import CRUDBase
from app.models.medical_record import MedicalRecord
from app.schemas.medical_record import MedicalRecordCreate, MedicalRecordUpdate


class CRUDMedicalRecord(CRUDBase[MedicalRecord, MedicalRecordCreate, MedicalRecordUpdate]):
    async def get_by_user(self, db: AsyncSession, *, user_id: UUID) -> List[MedicalRecord]:
        """获取用户的所有医疗档案"""
        result = await db.execute(
            select(self.model)
            .options(joinedload(self.model.file))
            .filter(self.model.user_id == user_id)
            .order_by(self.model.created_at.desc())
        )
        return result.scalars().all()
    
    async def get_by_file_id(self, db: AsyncSession, *, file_id: UUID) -> Optional[MedicalRecord]:
        """根据文件ID获取医疗档案"""
        result = await db.execute(
            select(self.model)
            .filter(self.model.file_id == file_id)
        )
        return result.scalar_one_or_none()
    
    async def get_user_records_paginated(
        self, 
        db: AsyncSession, 
        *, 
        user_id: UUID, 
        page: int = 1, 
        page_size: int = 20
    ) -> Tuple[List[MedicalRecord], int]:
        """获取用户医疗档案列表（分页）"""
        # Calculate offset
        offset = (page - 1) * page_size
        
        # Get total count
        count_result = await db.execute(
            select(func.count(self.model.id))
            .filter(self.model.user_id == user_id)
        )
        total = count_result.scalar() or 0
        
        # Get paginated results
        records_result = await db.execute(
            select(self.model)
            .options(joinedload(self.model.file))
            .filter(self.model.user_id == user_id)
            .order_by(self.model.upload_time.desc())
            .offset(offset)
            .limit(page_size)
        )
        records = records_result.scalars().all()
        
        return list(records), total

    async def update_extraction_status(
        self, 
        db: AsyncSession, 
        *, 
        record_id: UUID, 
        status: str,
        extracted_data: dict = None
    ) -> Optional[MedicalRecord]:
        """更新医疗信息提取状态和结果"""
        record = await self.get(db, id=record_id)
        if not record:
            return None
            
        # Update extraction status
        record.extraction_status = status
        
        # Update extracted medical information if provided
        if extracted_data:
            if "disease_name" in extracted_data:
                record.disease_name = extracted_data["disease_name"]
            if "report_category" in extracted_data:
                record.report_category = extracted_data["report_category"]
            if "report_name" in extracted_data:
                record.report_name = extracted_data["report_name"]
            if "report_date" in extracted_data:
                record.report_date = extracted_data["report_date"]
            if "hospital_name" in extracted_data:
                record.hospital_name = extracted_data["hospital_name"]
        
        await db.commit()
        await db.refresh(record)
        return record

    async def count_by_user(self, db: AsyncSession, *, user_id: UUID) -> int:
        """获取用户医疗档案数量"""
        result = await db.execute(
            select(func.count(self.model.id))
            .filter(self.model.user_id == user_id)
        )
        return result.scalar() or 0

    async def get_with_file(self, db: AsyncSession, *, id: UUID) -> Optional[MedicalRecord]:
        """获取医疗档案详情，包含关联的文件信息"""
        result = await db.execute(
            select(self.model)
            .options(joinedload(self.model.file))
            .filter(self.model.id == id)
        )
        return result.scalar_one_or_none()

    async def get_by_upload_time(
        self,
        db: AsyncSession,
        *,
        user_id: UUID,
        upload_time: datetime
    ) -> List[MedicalRecord]:
        """根据上传时间获取相同批次的医疗档案，包含文件信息"""
        result = await db.execute(
            select(self.model)
            .options(joinedload(self.model.file))
            .filter(
                self.model.user_id == user_id,
                self.model.upload_time == upload_time
            )
            .order_by(self.model.created_at.desc())
        )
        return result.scalars().all()


crud_medical_record = CRUDMedicalRecord(MedicalRecord)